import pandas as pd
from datetime import datetime
from models.models import db, Student, Parent, Class, Registration, Payment, Project, Account, Grade, Semester

class ImportHandler:
    def __init__(self):
        self.allowed_extensions = {'xlsx', 'xls'}
        self.required_sheets = {
            'grades': ['名称', '排序'],
            'semesters': ['名称', '排序'],
            'projects': ['名称', '排序'],
            'accounts': ['账号'],
            'parents': ['姓名', '电话'],
            'students': ['姓名', '性别', '出生日期', '身份证号', '地址', '家长姓名', '家长电话', '与家长关系'],
            'classes': ['年级', '年份', '学期'],
            'registrations': ['学生姓名', '班级年级', '班级年份', '班级学期'],
            'payments': ['日期', '学生姓名', '班级年级', '班级年份', '班级学期', '项目', '金额', '账号', '摘要']
        }

    def allowed_file(self, filename):
        return '.' in filename and filename.rsplit('.', 1)[1].lower() in self.allowed_extensions

    def import_data(self, file):
        try:
            # 读取Excel文件中的所有sheet
            excel_file = pd.ExcelFile(file)
            
            # 验证所需的sheet是否都存在
            missing_sheets = set(self.required_sheets.keys()) - set(excel_file.sheet_names)
            if missing_sheets:
                return False, f"Excel文件缺少以下sheet：{', '.join(missing_sheets)}"

            # 按照顺序导入数据（确保外键关系正确）
            import_order = ['grades', 'semesters', 'projects', 'accounts', 'parents', 'students', 'classes', 'registrations', 'payments']
            
            success_count = {}
            error_count = {}
            
            for sheet_name in import_order:
                df = pd.read_excel(file, sheet_name=sheet_name)
                
                # 验证必需列是否存在
                missing_columns = set(self.required_sheets[sheet_name]) - set(df.columns)
                if missing_columns:
                    return False, f"{sheet_name} sheet缺少以下列：{', '.join(missing_columns)}"
                
                # 导入数据
                success, error = getattr(self, f"_import_{sheet_name}")(df)
                success_count[sheet_name] = success
                error_count[sheet_name] = error

            # 提交所有更改
            try:
                db.session.commit()
                
                # 生成导入报告
                report = []
                for sheet_name in import_order:
                    report.append(f"{sheet_name}: 成功{success_count[sheet_name]}条，失败{error_count[sheet_name]}条")
                
                return True, "导入完成。\n" + "\n".join(report)
            except Exception as e:
                db.session.rollback()
                return False, f"导入失败：{str(e)}"
                
        except Exception as e:
            return False, f"导入失败：{str(e)}"

    def _import_grades(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                # 检查年级是否已存在
                existing_grade = Grade.query.filter_by(name=row['名称']).first()
                if existing_grade:
                    # 如果存在，更新排序
                    existing_grade.sort_order = int(row['排序'])
                    success_count += 1
                    continue
                    
                grade = Grade(
                    name=row['名称'],
                    sort_order=int(row['排序'])
                )
                db.session.add(grade)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_semesters(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                # 检查学期是否已存在
                existing_semester = Semester.query.filter_by(name=row['名称']).first()
                if existing_semester:
                    # 如果存在，更新排序
                    existing_semester.sort_order = int(row['排序'])
                    success_count += 1
                    continue
                    
                semester = Semester(
                    name=row['名称'],
                    sort_order=int(row['排序'])
                )
                db.session.add(semester)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_projects(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                # 检查项目是否已存在
                existing_project = Project.query.filter_by(name=row['名称']).first()
                if existing_project:
                    # 如果存在，更新排序
                    existing_project.sort_order = int(row['排序'])
                    success_count += 1
                    continue
                    
                project = Project(
                    name=row['名称'],
                    sort_order=int(row['排序'])
                )
                db.session.add(project)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_accounts(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                # 检查账号是否已存在
                existing_account = Account.query.filter_by(account=row['账号']).first()
                if existing_account:
                    # 如果账号已存在，跳过但不计入错误
                    continue
                    
                account = Account(
                    account=row['账号']
                )
                db.session.add(account)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_parents(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                parent = Parent(
                    name=row['姓名'],
                    phone=row['电话']
                )
                db.session.add(parent)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_students(self, df):
        success_count = 0
        error_count = 0

        # 先提交之前的更改以获取家长ID
        db.session.commit()

        for _, row in df.iterrows():
            try:
                # 查找或创建家长
                parent = Parent.query.filter_by(
                    name=row['家长姓名'],
                    phone=row['家长电话']
                ).first()
                
                if not parent:
                    parent = Parent(
                        name=row['家长姓名'],
                        phone=row['家长电话']
                    )
                    db.session.add(parent)
                    db.session.commit()  # 提交以获取parent.id

                student = Student(
                    name=row['姓名'],
                    gender=row['性别'],
                    birth_date=pd.to_datetime(row['出生日期']).date(),
                    id_card=row['身份证号'],
                    address=row['地址'],
                    parent_id=parent.id,
                    parent_name=parent.name,
                    phone=parent.phone,
                    parent_relation=row['与家长关系']
                )
                db.session.add(student)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_classes(self, df):
        success_count = 0
        error_count = 0

        for _, row in df.iterrows():
            try:
                class_obj = Class(
                    grade=row['年级'],
                    year=row['年份'],
                    semester=row['学期']
                )
                db.session.add(class_obj)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_registrations(self, df):
        success_count = 0
        error_count = 0

        # 先提交之前的更改以确保能查询到学生和班级
        db.session.commit()

        for _, row in df.iterrows():
            try:
                # 查找学生
                student = Student.query.filter_by(name=row['学生姓名']).first()
                if not student:
                    error_count += 1
                    continue

                # 查找班级
                class_obj = Class.query.filter_by(
                    grade=row['班级年级'],
                    year=row['班级年份'],
                    semester=row['班级学期']
                ).first()
                if not class_obj:
                    error_count += 1
                    continue

                registration = Registration(
                    student_id=student.id,
                    class_id=class_obj.id,
                    year=class_obj.year,
                    semester=class_obj.semester,
                    grade=class_obj.grade,
                    name=student.name,
                    gender=student.gender,
                    parent_relation=student.parent_relation,
                    student=student.name
                )
                db.session.add(registration)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count

    def _import_payments(self, df):
        success_count = 0
        error_count = 0

        # 先提交之前的更改以确保能查询到注册记录
        db.session.commit()

        for _, row in df.iterrows():
            try:
                # 查找注册记录
                registration = Registration.query.join(Student).join(Class).filter(
                    Student.name == row['学生姓名'],
                    Class.grade == row['班级年级'],
                    Class.year == row['班级年份'],
                    Class.semester == row['班级学期']
                ).first()

                if not registration:
                    error_count += 1
                    continue

                payment = Payment(
                    date=pd.to_datetime(row['日期']).date(),
                    registration_id=registration.id,
                    student=registration.student,
                    class_name=f"{registration.grade} {registration.year} {registration.semester}",
                    year=registration.year,
                    semester=registration.semester,
                    grade=registration.grade,
                    project=row['项目'],
                    amount=float(row['金额']),
                    account=row['账号'],
                    summary=row['摘要'] if pd.notna(row['摘要']) else None
                )
                db.session.add(payment)
                success_count += 1
            except Exception as e:
                error_count += 1
                continue

        return success_count, error_count 